Enable Union notation at Python-scope for warp arrays#1549
Enable Union notation at Python-scope for warp arrays#1549FabienPean-Virtonomy wants to merge 1 commit into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds ChangesPEP 604 Union Support for Array Annotations
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
warp/_src/types.py (1)
5132-5153:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReturning classes here breaks the existing
None-argument inference path.After this change,
wp.array[...]is a class object, butinfer_argument_types()still reconstructs array templates withtype(t)(dtype=t.dtype, ndim=t.ndim)on Line 7163. For these annotations,type(t)is now_ArrayAnnotationTypeMeta, so passingNonefor an array-typed parameter will raise instead of preserving the template type.At minimum, the downstream reconstruction needs to special-case
is_array_annotation(t)and reuset(or rebuild it via_make_array_annotation_type(...)) instead of callingtype(t)(...).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@warp/_src/types.py` around lines 5132 - 5153, The change made in _parse_array_subscript returns class-like objects which breaks infer_argument_types' reconstruction that calls type(t)(dtype=..., ndim=...) (type(t) is now _ArrayAnnotationTypeMeta); update infer_argument_types to special-case array annotations by using is_array_annotation(t) and either reuse the existing annotation object t directly when filling in None arguments or rebuild the annotation via _make_array_annotation_type(_ARRAY_ANNOTATION_MAP.get(base), dtype=t.dtype, ndim=t.ndim) instead of calling type(t)(...), ensuring array templates preserve their original semantics.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@warp/_src/types.py`:
- Around line 5112-5129: _make_array_annotation_type currently constructs a new
_ArrayAnnotationTypeMeta on every call causing identity instability; add a cache
(e.g., a module-level dict or an attribute on ann_base) keyed by the normalized
dtype and ndim (use the same dtype = dtype if dtype is Any else
type_to_warp(dtype) logic and the resolved ndim) and return the cached class
when present instead of allocating a new one; update the function to compute the
cache key (ann_base, dtype, ndim), check and return from cache, and only create
and store a new _ArrayAnnotationTypeMeta if the key is missing so repeated
subscripts (e.g. wp.array[float]) yield identical types.
- Around line 5064-5070: The __ctype__ implementation for array-like classes
currently returns a bare array_t() which discards annotated rank/shape/strides;
change __ctype__ (using cls._concrete_cls, array_t(), indexedarray_t,
ARRAY_MAX_DIMS) to preserve the annotated ndim/shape/strides the same way the
old annotation path did—read cls.ndim (falling back to Any), cls.shape and
cls.strides if present (or default placeholders) and construct array_t(...) with
those values instead of returning array_t() so codegen/native consumers retain
the array rank information.
---
Outside diff comments:
In `@warp/_src/types.py`:
- Around line 5132-5153: The change made in _parse_array_subscript returns
class-like objects which breaks infer_argument_types' reconstruction that calls
type(t)(dtype=..., ndim=...) (type(t) is now _ArrayAnnotationTypeMeta); update
infer_argument_types to special-case array annotations by using
is_array_annotation(t) and either reuse the existing annotation object t
directly when filling in None arguments or rebuild the annotation via
_make_array_annotation_type(_ARRAY_ANNOTATION_MAP.get(base), dtype=t.dtype,
ndim=t.ndim) instead of calling type(t)(...), ensuring array templates preserve
their original semantics.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Enterprise
Run ID: 25e54390-be3e-49bd-94f0-acb5c563f40b
📒 Files selected for processing (3)
CHANGELOG.mdwarp/_src/types.pywarp/tests/test_subscript_types.py
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def __eq__(cls, other): | ||
| if not isinstance(other, type): | ||
| return NotImplemented | ||
| if not issubclass(other, _ArrayAnnotationBase): | ||
| return False |
There was a problem hiding this comment.
__eq__ implementation returns False (not NotImplemented) when other is a type but is not an _ArrayAnnotationBase subclass. Python convention is to return NotImplemented when a comparison is not meaningful for the operand types, so the other side can still be tried. Returning False suppresses that fallback. In practice this rarely matters here (built-in type.__eq__ also falls back correctly), but it diverges from the standard contract and could silently short-circuit user-defined __eq__ on custom type objects compared against annotation types.
| def __eq__(cls, other): | |
| if not isinstance(other, type): | |
| return NotImplemented | |
| if not issubclass(other, _ArrayAnnotationBase): | |
| return False | |
| def __eq__(cls, other): | |
| if not isinstance(other, type): | |
| return NotImplemented | |
| if not issubclass(other, _ArrayAnnotationBase): | |
| return NotImplemented |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
|
/ok to test 2b76761 |
|
Thanks for the fix. It works and the tests are thorough. The root cause is just that |
FabienPean-Virtonomy
force-pushed
the
FabienPean/union-type-hint
branch
from
3d8266f to
379fb1e
Compare
|
Two factors that led to this version:
Your option leads indeed to a much smaller fix. The differences are that |
FabienPean-Virtonomy
changed the title
|
Thanks for the fix! It looks good. Could you please rebase your changes on We can revisit whether these annotations should be real type objects in the future if the current approach runs into a concrete limitation, but I think the simpler approach addresses the immediate need. |
…d tests (NVIDIAGH-1548) Signed-off-by: Fabien Péan <pean@virtonomy.io>
FabienPean-Virtonomy
force-pushed
the
FabienPean/union-type-hint
branch
from
379fb1e to
565f925
Compare
2 participants
Description
Array subscript annotations such as wp.array[...] still return lightweight annotation instances but they now support PEP 604 union expressions at Python scope through instance operator support (for example,
wp.array[float] | float and float | wp.array[float]).Closes #1548
Checklist
Validation Summary
uv run test_subscript_types.py
Bug Fix Repro
Summary by CodeRabbit
Release Notes
Bug Fixes
|) with array type annotations, enabling syntax likewp.array[float] | floatin Python method contexts.Tests